import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QLineEdit, QPushButton, QFileDialog, QTextEdit,
                             QComboBox, QTabWidget, QMessageBox, QGroupBox, QRadioButton)
from PyQt5.QtCore import Qt
from custom_crypto_system import crypto_system

class CryptoGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("自定义密码加密系统")
        self.setGeometry(100, 100, 800, 600)
        
        # 创建主窗口部件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        
        # 创建主布局
        self.main_layout = QVBoxLayout()
        self.central_widget.setLayout(self.main_layout)
        
        # 创建选项卡
        self.tab_widget = QTabWidget()
        self.main_layout.addWidget(self.tab_widget)
        
        # 创建各个选项卡
        self.create_password_crypto_tab()  # 新增密码加密选项卡
        self.create_rsa_tab()
        self.create_keygen_tab()
        
    def create_password_crypto_tab(self):
        """创建基于密码的加密/解密选项卡"""
        pw_widget = QWidget()
        pw_layout = QVBoxLayout()
        pw_widget.setLayout(pw_layout)
        
        # 密码加密组
        encrypt_group = QGroupBox("基于密码的文件加密")
        encrypt_layout = QVBoxLayout()
        encrypt_group.setLayout(encrypt_layout)
        
        # 文件选择
        file_layout = QHBoxLayout()
        self.pw_encrypt_file_path = QLineEdit()
        self.pw_encrypt_file_path.setPlaceholderText("选择要加密的文件")
        encrypt_file_btn = QPushButton("浏览")
        encrypt_file_btn.clicked.connect(self.browse_pw_encrypt_file)
        file_layout.addWidget(QLabel("文件:"))
        file_layout.addWidget(self.pw_encrypt_file_path)
        file_layout.addWidget(encrypt_file_btn)
        encrypt_layout.addLayout(file_layout)
        
        # 密码输入
        password_layout = QHBoxLayout()
        self.pw_encrypt_password = QLineEdit()
        self.pw_encrypt_password.setPlaceholderText("输入加密密码")
        self.pw_encrypt_password.setEchoMode(QLineEdit.Password)
        password_layout.addWidget(QLabel("密码:"))
        password_layout.addWidget(self.pw_encrypt_password)
        encrypt_layout.addLayout(password_layout)
        
        # 确认密码
        confirm_layout = QHBoxLayout()
        self.pw_encrypt_confirm = QLineEdit()
        self.pw_encrypt_confirm.setPlaceholderText("确认密码")
        self.pw_encrypt_confirm.setEchoMode(QLineEdit.Password)
        confirm_layout.addWidget(QLabel("确认:"))
        confirm_layout.addWidget(self.pw_encrypt_confirm)
        encrypt_layout.addLayout(confirm_layout)
        
        # 加密按钮
        encrypt_btn = QPushButton("加密文件")
        encrypt_btn.clicked.connect(self.password_encrypt_file)
        encrypt_layout.addWidget(encrypt_btn)
        
        pw_layout.addWidget(encrypt_group)
        
        # 密码解密组
        decrypt_group = QGroupBox("基于密码的文件解密")
        decrypt_layout = QVBoxLayout()
        decrypt_group.setLayout(decrypt_layout)
        
        # 加密文件选择
        enc_file_layout = QHBoxLayout()
        self.pw_decrypt_file_path = QLineEdit()
        self.pw_decrypt_file_path.setPlaceholderText("选择要解密的文件")
        decrypt_file_btn = QPushButton("浏览")
        decrypt_file_btn.clicked.connect(self.browse_pw_decrypt_file)
        enc_file_layout.addWidget(QLabel("文件:"))
        enc_file_layout.addWidget(self.pw_decrypt_file_path)
        enc_file_layout.addWidget(decrypt_file_btn)
        decrypt_layout.addLayout(enc_file_layout)
        
        # 密码输入
        pw_decrypt_layout = QHBoxLayout()
        self.pw_decrypt_password = QLineEdit()
        self.pw_decrypt_password.setPlaceholderText("输入解密密码")
        self.pw_decrypt_password.setEchoMode(QLineEdit.Password)
        pw_decrypt_layout.addWidget(QLabel("密码:"))
        pw_decrypt_layout.addWidget(self.pw_decrypt_password)
        decrypt_layout.addLayout(pw_decrypt_layout)
        
        # 解密按钮
        decrypt_btn = QPushButton("解密文件")
        decrypt_btn.clicked.connect(self.password_decrypt_file)
        decrypt_layout.addWidget(decrypt_btn)
        
        pw_layout.addWidget(decrypt_group)
        
        # 结果显示
        self.pw_result = QTextEdit()
        self.pw_result.setReadOnly(True)
        self.pw_result.setPlaceholderText("操作结果将显示在这里...")
        pw_layout.addWidget(QLabel("操作结果:"))
        pw_layout.addWidget(self.pw_result)
        
        self.tab_widget.addTab(pw_widget, "密码加密")
    
    def create_rsa_tab(self):
        """创建RSA加密/解密选项卡"""
        rsa_widget = QWidget()
        rsa_layout = QVBoxLayout()
        rsa_widget.setLayout(rsa_layout)
        
        # RSA加密组
        encrypt_group = QGroupBox("RSA文件加密")
        encrypt_layout = QVBoxLayout()
        encrypt_group.setLayout(encrypt_layout)
        
        # 文件选择
        file_layout = QHBoxLayout()
        self.encrypt_file_path = QLineEdit()
        self.encrypt_file_path.setPlaceholderText("选择要加密的文件")
        encrypt_file_btn = QPushButton("浏览")
        encrypt_file_btn.clicked.connect(self.browse_encrypt_file)
        file_layout.addWidget(QLabel("文件:"))
        file_layout.addWidget(self.encrypt_file_path)
        file_layout.addWidget(encrypt_file_btn)
        encrypt_layout.addLayout(file_layout)
        
        # 公钥选择
        pub_key_layout = QHBoxLayout()
        self.pub_key_path = QLineEdit()
        self.pub_key_path.setPlaceholderText("选择公钥文件")
        pub_key_btn = QPushButton("浏览")
        pub_key_btn.clicked.connect(self.browse_pub_key)
        pub_key_layout.addWidget(QLabel("公钥:"))
        pub_key_layout.addWidget(self.pub_key_path)
        pub_key_layout.addWidget(pub_key_btn)
        encrypt_layout.addLayout(pub_key_layout)
        
        # 加密按钮
        encrypt_btn = QPushButton("加密文件")
        encrypt_btn.clicked.connect(self.encrypt_file)
        encrypt_layout.addWidget(encrypt_btn)
        
        rsa_layout.addWidget(encrypt_group)
        
        # RSA解密组
        decrypt_group = QGroupBox("RSA文件解密")
        decrypt_layout = QVBoxLayout()
        decrypt_group.setLayout(decrypt_layout)
        
        # 加密文件选择
        enc_file_layout = QHBoxLayout()
        self.decrypt_file_path = QLineEdit()
        self.decrypt_file_path.setPlaceholderText("选择要解密的文件")
        decrypt_file_btn = QPushButton("浏览")
        decrypt_file_btn.clicked.connect(self.browse_decrypt_file)
        enc_file_layout.addWidget(QLabel("文件:"))
        enc_file_layout.addWidget(self.decrypt_file_path)
        enc_file_layout.addWidget(decrypt_file_btn)
        decrypt_layout.addLayout(enc_file_layout)
        
        # 私钥选择
        priv_key_layout = QHBoxLayout()
        self.priv_key_path = QLineEdit()
        self.priv_key_path.setPlaceholderText("选择私钥文件")
        priv_key_btn = QPushButton("浏览")
        priv_key_btn.clicked.connect(self.browse_priv_key)
        priv_key_layout.addWidget(QLabel("私钥:"))
        priv_key_layout.addWidget(self.priv_key_path)
        priv_key_layout.addWidget(priv_key_btn)
        decrypt_layout.addLayout(priv_key_layout)
        
        # 密码输入
        password_layout = QHBoxLayout()
        self.password_input = QLineEdit()
        self.password_input.setPlaceholderText("输入私钥密码(如果有)")
        self.password_input.setEchoMode(QLineEdit.Password)
        password_layout.addWidget(QLabel("密码:"))
        password_layout.addWidget(self.password_input)
        decrypt_layout.addLayout(password_layout)
        
        # 解密按钮
        decrypt_btn = QPushButton("解密文件")
        decrypt_btn.clicked.connect(self.decrypt_file)
        decrypt_layout.addWidget(decrypt_btn)
        
        rsa_layout.addWidget(decrypt_group)
        
        # 结果显示
        self.rsa_result = QTextEdit()
        self.rsa_result.setReadOnly(True)
        self.rsa_result.setPlaceholderText("操作结果将显示在这里...")
        rsa_layout.addWidget(QLabel("操作结果:"))
        rsa_layout.addWidget(self.rsa_result)
        
        self.tab_widget.addTab(rsa_widget, "RSA加密")
    
    def create_keygen_tab(self):
        """创建密钥生成选项卡"""
        keygen_widget = QWidget()
        keygen_layout = QVBoxLayout()
        keygen_widget.setLayout(keygen_layout)
        
        # 密钥长度选择
        key_size_layout = QHBoxLayout()
        key_size_layout.addWidget(QLabel("密钥长度:"))
        self.key_size = QComboBox()
        self.key_size.addItems(["1024", "2048", "4096"])
        self.key_size.setCurrentText("2048")
        key_size_layout.addWidget(self.key_size)
        key_size_layout.addStretch()
        keygen_layout.addLayout(key_size_layout)
        
        # 保存路径
        save_path_layout = QHBoxLayout()
        self.key_save_path = QLineEdit()
        self.key_save_path.setPlaceholderText("选择密钥保存路径")
        save_path_btn = QPushButton("浏览")
        save_path_btn.clicked.connect(self.browse_key_save_path)
        save_path_layout.addWidget(QLabel("保存路径:"))
        save_path_layout.addWidget(self.key_save_path)
        save_path_layout.addWidget(save_path_btn)
        keygen_layout.addLayout(save_path_layout)
        
        # 是否使用密码保护私钥
        password_protect_layout = QHBoxLayout()
        self.use_password_check = QRadioButton("使用密码保护私钥")
        self.no_password_check = QRadioButton("不使用密码保护")
        self.no_password_check.setChecked(True)
        
        password_protect_layout.addWidget(self.use_password_check)
        password_protect_layout.addWidget(self.no_password_check)
        password_protect_layout.addStretch()
        keygen_layout.addLayout(password_protect_layout)
        
        # 密码输入
        self.priv_key_password_layout = QHBoxLayout()
        self.priv_key_password = QLineEdit()
        self.priv_key_password.setPlaceholderText("输入私钥保护密码")
        self.priv_key_password.setEchoMode(QLineEdit.Password)
        self.priv_key_password_layout.addWidget(QLabel("密码:"))
        self.priv_key_password_layout.addWidget(self.priv_key_password)
        keygen_layout.addLayout(self.priv_key_password_layout)
        self.priv_key_password_layout.setEnabled(False)
        
        # 确认密码
        self.priv_key_confirm_layout = QHBoxLayout()
        self.priv_key_confirm = QLineEdit()
        self.priv_key_confirm.setPlaceholderText("确认密码")
        self.priv_key_confirm.setEchoMode(QLineEdit.Password)
        self.priv_key_confirm_layout.addWidget(QLabel("确认:"))
        self.priv_key_confirm_layout.addWidget(self.priv_key_confirm)
        keygen_layout.addLayout(self.priv_key_confirm_layout)
        self.priv_key_confirm_layout.setEnabled(False)
        
        # 连接信号
        self.use_password_check.toggled.connect(self.toggle_password_fields)
        
        # 生成按钮
        generate_btn = QPushButton("生成密钥对")
        generate_btn.clicked.connect(self.generate_keys)
        keygen_layout.addWidget(generate_btn)
        
        # 结果显示
        self.keygen_result = QTextEdit()
        self.keygen_result.setReadOnly(True)
        self.keygen_result.setPlaceholderText("密钥生成结果将显示在这里...")
        keygen_layout.addWidget(QLabel("生成结果:"))
        keygen_layout.addWidget(self.keygen_result)
        
        keygen_layout.addStretch()
        
        self.tab_widget.addTab(keygen_widget, "密钥生成")
    
    def toggle_password_fields(self):
        """切换密码输入字段的可用性"""
        use_password = self.use_password_check.isChecked()
        self.priv_key_password_layout.setEnabled(use_password)
        self.priv_key_confirm_layout.setEnabled(use_password)
    
    # 浏览文件的方法
    def browse_pw_encrypt_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择要加密的文件")
        if file_path:
            self.pw_encrypt_file_path.setText(file_path)
    
    def browse_pw_decrypt_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择要解密的文件", "", "加密文件 (*.encrypted)")
        if file_path:
            self.pw_decrypt_file_path.setText(file_path)
    
    def browse_encrypt_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择要加密的文件")
        if file_path:
            self.encrypt_file_path.setText(file_path)
    
    def browse_decrypt_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择要解密的文件", "", "RSA加密文件 (*.rsa_encrypted)")
        if file_path:
            self.decrypt_file_path.setText(file_path)
    
    def browse_pub_key(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择公钥文件", "", "PEM文件 (*.pem)")
        if file_path:
            self.pub_key_path.setText(file_path)
    
    def browse_priv_key(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择私钥文件", "", "PEM文件 (*.pem)")
        if file_path:
            self.priv_key_path.setText(file_path)
    
    def browse_key_save_path(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择密钥保存目录")
        if dir_path:
            self.key_save_path.setText(dir_path)
    
    # 密码加密/解密方法
    def password_encrypt_file(self):
        """使用密码加密文件"""
        try:
            file_path = self.pw_encrypt_file_path.text()
            password = self.pw_encrypt_password.text()
            confirm_password = self.pw_encrypt_confirm.text()
            
            if not file_path or not password:
                QMessageBox.warning(self, "警告", "请选择要加密的文件并输入密码")
                return
            
            if password != confirm_password:
                QMessageBox.warning(self, "警告", "两次输入的密码不一致")
                return
            
            # 加密文件
            output_path = crypto_system.encrypt_file_with_password(file_path, password)
            
            self.pw_result.setText(f"文件加密成功！\n加密文件保存为: {output_path}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))
    
    def password_decrypt_file(self):
        """使用密码解密文件"""
        try:
            file_path = self.pw_decrypt_file_path.text()
            password = self.pw_decrypt_password.text()
            
            if not file_path or not password:
                QMessageBox.warning(self, "警告", "请选择要解密的文件并输入密码")
                return
            
            # 解密文件
            output_path = crypto_system.decrypt_file_with_password(file_path, password)
            
            self.pw_result.setText(f"文件解密成功！\n解密文件保存为: {output_path}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))
    
    # RSA加密/解密方法
    def encrypt_file(self):
        """使用RSA加密文件"""
        try:
            file_path = self.encrypt_file_path.text()
            pub_key_path = self.pub_key_path.text()
            
            if not file_path or not pub_key_path:
                QMessageBox.warning(self, "警告", "请选择要加密的文件和公钥文件")
                return
            
            # 加密文件
            output_path = crypto_system.encrypt_file_with_rsa(file_path, pub_key_path)
            
            self.rsa_result.setText(f"文件加密成功！\n加密文件保存为: {output_path}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))
    
    def decrypt_file(self):
        """使用RSA解密文件"""
        try:
            file_path = self.decrypt_file_path.text()
            priv_key_path = self.priv_key_path.text()
            password = self.password_input.text()
            
            if not file_path or not priv_key_path:
                QMessageBox.warning(self, "警告", "请选择要解密的文件和私钥文件")
                return
            
            # 解密文件
            output_path = crypto_system.decrypt_file_with_rsa(
                file_path, priv_key_path, password if password else None
            )
            
            self.rsa_result.setText(f"文件解密成功！\n解密文件保存为: {output_path}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))
    
    # 密钥生成方法
    def generate_keys(self):
        """生成RSA密钥对"""
        try:
            key_size = int(self.key_size.currentText())
            save_path = self.key_save_path.text()
            use_password = self.use_password_check.isChecked()
            
            if not save_path:
                QMessageBox.warning(self, "警告", "请选择密钥保存路径")
                return
            
            # 检查密码
            password = None
            if use_password:
                password = self.priv_key_password.text()
                confirm_password = self.priv_key_confirm.text()
                
                if not password:
                    QMessageBox.warning(self, "警告", "请输入私钥保护密码")
                    return
                
                if password != confirm_password:
                    QMessageBox.warning(self, "警告", "两次输入的密码不一致")
                    return
            
            # 生成密钥对
            public_key, private_key = crypto_system.generate_rsa_keys_with_password(key_size)
            
            # 保存密钥对
            pub_key_path = os.path.join(save_path, "public_key.pem")
            priv_key_path = os.path.join(save_path, "private_key.pem")
            
            crypto_system.save_rsa_keys(
                public_key, private_key, pub_key_path, priv_key_path, 
                password if use_password else None
            )
            
            result_text = f"密钥对生成成功！\n密钥长度: {key_size}位\n公钥保存为: {pub_key_path}\n私钥保存为: {priv_key_path}"
            if use_password:
                result_text += "\n(私钥已使用密码保护)"
            
            self.keygen_result.setText(result_text)
            
        except Exception as e:
            QMessageBox.critical(self, "错误", str(e))

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CryptoGUI()
    window.show()
    sys.exit(app.exec_())